#include "js_ColorCube.h"

using namespace DD::Image;

// helper functions

float js_ColorCube::deg_to_rad(float deg)
{
    return deg * M_PI / 180.0f;
}

void js_ColorCube::drawCircle(const float radius, const float r, const float g, const float b)
{
    float x, y;
    glBegin(GL_LINES);
    glColor3f(r, g, b);
    x = (float) radius * cos(deg_to_rad(350));
    y = (float) radius * sin(deg_to_rad(350));
    for (int j = 0; j < 360;)
    {
        glVertex2f(x, y);
        x = (float) radius * cos(deg_to_rad(j));
        y = (float) radius * sin(deg_to_rad(j));
        glVertex2f(x, y);
        j += 10;
    }
    glEnd();
}

void js_ColorCube::drawBox(const float size, const float r, const float g, const float b, const bool color )
{
    if ( !color )
        glColor3f( 1.0f, 1.0f, 1.0f );

    glBegin( GL_LINES );
    {
        if ( color )
            glColor3f( 0.0f, 0.0f, 0.0f );
        glVertex3f( 0.0f, 0.0f, 0.0f );
        if ( color )
            glColor3f( size, 0.0f, 0.0f );
        glVertex3f( size, 0.0f, 0.0f );

        if ( color )
            glColor3f( 0.0f, 0.0f, 0.0f );
        glVertex3f( 0.0f, 0.0f, 0.0f );
        if ( color )
            glColor3f( 0.0f, size, 0.0f );
        glVertex3f( 0.0f, size, 0.0f );

        if ( color )
            glColor3f( size, 0.0f, 0.0f );
        glVertex3f( size, 0.0f, 0.0f );
        if ( color )
            glColor3f( size, size, 0.0f );
        glVertex3f( size, size, 0.0f );
        
        if ( color )
            glColor3f( 0.0f, size, 0.0f );
        glVertex3f( 0.0f, size, 0.0f );
        if ( color )
            glColor3f( size, size, 0.0f );
        glVertex3f( size, size, 0.0f );

        if ( color )
            glColor3f( 0.0f, 0.0f, size );
        glVertex3f( 0.0f, 0.0f, size );
        if ( color )
            glColor3f( size, 0.0f, size );
        glVertex3f( size, 0.0f, size );

        if ( color )
            glColor3f( 0.0f, 0.0f, size );
        glVertex3f( 0.0f, 0.0f, size );
        if ( color )
            glColor3f( 0.0f, size, size );
        glVertex3f( 0.0f, size, size );

        if ( color )
            glColor3f( size, 0.0f, size );
        glVertex3f( size, 0.0f, size );
        if ( color )
            glColor3f( size, size, size );
        glVertex3f( size, size, size );
        
        if ( color )
            glColor3f( 0.0f, size, size );
        glVertex3f( 0.0f, size, size );
        if ( color )
            glColor3f( size, size, size );
        glVertex3f( size, size, size );

        if ( color )
            glColor3f( 0.0f, 0.0f, 0.0f );
        glVertex3f( 0.0f, 0.0f, 0.0f );
        if ( color )
            glColor3f( 0.0f, 0.0f, size );
        glVertex3f( 0.0f, 0.0f, size );

        if ( color )
            glColor3f( 0.0f, size, 0.0f );
        glVertex3f( 0.0f, size, 0.0f );
        if ( color )
            glColor3f( 0.0f, size, size );
        glVertex3f( 0.0f, size, size );

        if ( color )
            glColor3f( size, 0.0f, 0.0f );
        glVertex3f( size, 0.0f, 0.0f );
        if ( color )
            glColor3f( size, 0.0f, size );
        glVertex3f( size, 0.0f, size );

        if ( color )
            glColor3f( size, size, 0.0f );
        glVertex3f( size, size, 0.0f );
        if ( color )
            glColor3f( size, size, size );
        glVertex3f( size, size, size );
    }
    glEnd();
}

float js_ColorCube::luma( const float r, const float g, const float b )
{
    return 0.2125f * r + 0.7154f * g + 0.0721f * b;
}

void js_ColorCube::updateMatrices( void )
{
    AxisOp * _tol = dynamic_cast<AxisOp*>(Op::input(1));
    AxisOp * _soft = dynamic_cast<AxisOp*>(Op::input(2));

    if ( _tol )
    {
        _tol->validate();
        m_tolMat = _tol->matrix();
        m_tolMatInv = _tol->imatrix();
    }

    if ( _soft )
    {
        _soft->validate();
        m_softMat = _soft->matrix();
        m_softMatInv = _soft->imatrix();
    }
}

float js_ColorCube::js_3dKey( const Vector3 & pos )
{
    Vector3 _posTol = m_tolMatInv.transform( pos ); // position in relation to the Tolerance sphere
    Vector3 _posSoft = m_softMatInv.transform( pos ); // position in relation to the Softness sphere

    // inside the tolerance meaning we key it away!
    float result  = 0.0f;
    float _posTolLen = _posTol.length();
    if ( _posTolLen > 1.0f )
    {
        float _posSoftLen = _posSoft.length();

        if ( _posSoftLen > 1.0f )
        {
            // outside the softness ( and tolerance ) resulting in a white alpha
            result = 1.0f;
        }
        else
        {
            // the position of the color on a unit sphere
            // transformed back to the colorspace
            Vector3 _posTolNorm = m_tolMat.transform( _posTol / _posTolLen );
            Vector3 _posSoftNorm = m_softMat.transform( _posSoft / _posSoftLen );

            // difference vectors
            Vector3 _distVec = _posTolNorm - _posSoftNorm;
            Vector3 _posVec = _posTolNorm - pos;

            // the resulting alpha
            result = _posVec.length() / _distVec.length();
        }
    }

    return result;
}

float js_ColorCube::js_KKey( const Vector3 & pos )
{
    float _min = std::numeric_limits<float>::max();

    std::vector<Vector3>::const_iterator it = m_keySample.begin();
    std::vector<Vector3>::const_iterator itEnd = m_keySample.end();

    for ( ; it != itEnd; ++it )
    {
        Vector3 _distVec = (*it) - pos;
        float _dist = _distVec.length();

        if ( _dist < _min )
            _min = _dist;
    }

    float result = 0.0f;

    if ( _min > m_kTolMin )
    {
        if ( _min > m_kTolMax )
        {
            result = 1.0f;
        }
        else
        {
            result = ( _min - m_kTolMin ) / ( m_kTolMax - m_kTolMin );
        }
    }

    return result;
}


void js_ColorCube::js_Warp( float & r, float & g, float & b )
{
    const float _luma = luma( r, g, b );
    Vector3 _pos( r, g, b );
    Vector3 _offset( 0.0f, 0.0f, 0.0f );
    Vector3 _r( _luma, 0.0f, 0.0f );
    Vector3 _g( 0.0f, _luma, 0.0f );
    Vector3 _b( 0.0f, 0.0f, _luma );
    Vector3 _sat( _luma, _luma, _luma );

    _offset += ( _pos - _r ) * m_warperCurve.getValue( 0, r );
    _offset += ( _pos - _g ) * m_warperCurve.getValue( 1, g );
    _offset += ( _pos - _b ) * m_warperCurve.getValue( 2, b );
    _offset += ( _pos - _sat ) * m_warperCurve.getValue( 3, _luma );

    r += _offset.x;
    g += _offset.y;
    b += _offset.z; 
}


js_ColorCube::js_ColorCube(Node *node)
    : Iop(node)
    , m_displayListId(-1)
    , m_drawSize(1.0f)
    , m_displaySample(false)
    , m_doKey(false)
    , m_doKeyDraw(false)
    , m_doDespill(false)
    , m_blueScreen(false)
    , m_colorBalance(0.5f)
    , m_colorMult(1.0f)
    , m_doWarp(false)
    , m_warperCurve(warperDefaults)
    , m_doKKey(false)
    , m_kTolMin(0.1f)
    , m_kTolMax(0.2f)
{
    inputs(3);
    m_centerColor[0] = 0.0f;
    m_centerColor[1] = 0.0f;
    m_centerColor[2] = 0.0f;

    m_sampleColor[0] = 0.0f;
    m_sampleColor[1] = 0.0f;
    m_sampleColor[2] = 0.0f;

    m_hash.reset();
    m_sampleHash.reset();
    m_tolMat.makeIdentity();
    m_tolMatInv.makeIdentity();
    m_softMat.makeIdentity();
    m_softMatInv.makeIdentity();
}

const char* js_ColorCube::input_label(int n, char*) const
{
    switch (n)
    {
        case 0:
            return "img";
        case 1:
            return "tol";
        case 2:
            return "soft";
    }
    return "";
}

Op* js_ColorCube::default_input(int n) const
{
    if ( n == 0 )
        return DD::Image::Iop::default_input(n);

    return DD::Image::Op::default_input(n);
}

void js_ColorCube::_validate(bool for_real)
{
    Iop::validate(for_real);
    Op::_validate(for_real); // mainly to get the axis nodes updated

    copy_info();
    info_.channels(Mask_RGBA);
    set_out_channels(Mask_RGBA);

    input0().validate(for_real);

    updateMatrices();
}
void js_ColorCube::_request(int x, int y, int r, int t, ChannelMask channels, int count)
{
    input0().request(x, y, r, t, Mask_RGBA, count);
}

bool js_ColorCube::test_input(int input, Op* op) const
{
    if (input == 0)
        if (dynamic_cast<Iop*> (op) != 0)
            return true;
    if (input == 1)
        if (dynamic_cast<AxisOp*> (op) != 0)
            return true;
    if (input == 2)
        if (dynamic_cast<AxisOp*> (op) != 0)
            return true;

    return false;
}

void js_ColorCube::engine(int y, int x, int r, ChannelMask channels, Row& out)
{
    Row _input(x, r);
    bool _blackit = true;

    if ( m_doKey || m_doWarp || m_doDespill )
    {
        _blackit = false;
        input0().get(y, x, r, Mask_RGBA, _input);
        const float * _r = _input[Chan_Red] + x;
        const float * _g = _input[Chan_Green] + x;
        const float * _b = _input[Chan_Blue] + x;
        const float * _a = _input[Chan_Alpha] + x;

        float *rOut = out.writable(Chan_Red) + x;
        float *gOut = out.writable(Chan_Green) + x;
        float *bOut = out.writable(Chan_Blue) + x;
        float *aOut = out.writable(Chan_Alpha) + x;

        const float *END = rOut + (r - x);
        const Vector3 _tol = m_tolMat.translation();

        while (rOut < END)
        {

            const Vector3 _pos( *_r, *_g, *_b );
            float R, G, B, A;
            R = *_r;
            G = *_g;
            B = *_b;
            A = *_a;

            if ( m_doKey )
                A = js_3dKey( _pos );

            if ( m_doDespill )
            {
                float _blend, _luma;

                if ( !m_blueScreen )
                {
                    _blend = R * m_colorBalance + B * ( 1.0f - m_colorBalance );
                    _blend *= m_colorMult;

                    G = std::min( G, _blend );
                    _luma = luma( R, G, B );
                    
                    R += _luma;
                    G += _luma;
                    B += _luma;
                }
                else
                {
                    _blend = R * m_colorBalance + G * ( 1.0f - m_colorBalance );
                    _blend *= m_colorMult;

                    B = std::min( B, _blend );
                    _luma = luma( R, G, B );
                    
                    R += _luma;
                    G += _luma;
                    B += _luma;
                }
            }

            if ( m_doWarp )
                js_Warp( R, G, B );

            *rOut = R;
            *gOut = G;
            *bOut = B;
            *aOut = A;

            ++rOut;
            ++gOut;
            ++bOut;
            ++aOut;
            ++_r;
            ++_g;
            ++_b;
        }
    }
    else
    {
        input0().get(y, x, r, Mask_All, _input);
        out.copy( _input, Mask_All, x, r ); 
    }

}

void js_ColorCube::updateColorList( void )
{
    Iop * _image = dynamic_cast<Iop*>(Op::input0());

    if ( _image )
    {
        _image->validate(true);
        int _imageWidth = _image->format().width();
        int _imageHeight = _image->format().height();
        _image->request( 0, 0, _imageWidth, _imageHeight, Mask_RGBA, 1 );
    
        if ( m_displayListId != -1 ) 
            glDeleteLists( m_displayListId, 1 );
        
        m_displayListId = glGenLists(1);

        // create display list
        glNewList( m_displayListId,GL_COMPILE );
        {
            // the color a points
            glBegin( GL_POINTS );
            {
                const Vector3 _tol = m_tolMat.translation();
                for( int y = 0; y < _imageHeight; ++y ) // loop over the lines
                {
                    Row _row( 0, _imageWidth - 1 );
                    _image->get( y, 0, _imageWidth - 1, Mask_RGBA, _row );

                    for( int x = 0; x < _imageWidth; ++x ) // loop over the colomns
                    {
                        const float * _r = _row[Chan_Red] + x;
                        const float * _g = _row[Chan_Green] + x;
                        const float * _b = _row[Chan_Blue] + x;

                        float R, G, B;

                        R = *_r;
                        G = *_g;
                        B = *_b;
                        Vector3 _pos( R, G, B );           

                        if ( m_doKeyDraw )
                        {
                            float _alpha = js_3dKey( _pos );

                            R = G = B = _alpha;
                        }
                        else
                        {
                            if ( m_doWarp )
                                js_Warp( R, G, B );
                        }
                        
                        glColor3f( R, G, B );

                        if ( !m_doWarp )
                            glVertex3f( _pos.x, _pos.y, _pos.z );
                        else
                            glVertex3f( R, G, B );

                        x+=4;
                    }
                }
            }
            glEnd();

            // white box around the colorspace
            drawBox( 1.0f, 1.0f, 1.0f, 1.0f, true );

            // keyer controls
            if ( m_doKey )
            {
                // tolerance
                glPushMatrix();
                glMultMatrixf(m_tolMat.array());
                drawCircle(1.0f, 1.0f, 1.0f, 1.0f );
                glRotatef(90.0f, 0.0f, 1.0f, 0.0f);
                drawCircle(1.0f, 1.0f, 1.0f, 1.0f );
                glPopMatrix();

                // softness
                glPushMatrix();
                glMultMatrixf(m_softMat.array());
                drawCircle(1.0f, 0.5f, 0.5f, 0.5f );
                glRotatef(90.0f, 0.0f, 1.0f, 0.0f);
                drawCircle(1.0f, 0.5f, 0.5f, 0.5f );
                glPopMatrix();
            }

            // sample a color
            if ( m_displaySample )
            {
                glPushMatrix();
                glTranslatef( m_sampleColor[0], m_sampleColor[1], m_sampleColor[2] );
                drawBox( 0.01f, m_sampleColor[0], m_sampleColor[1], m_sampleColor[2] );
                glPopMatrix();
            }
        }
        glEndList();
    }
}

void js_ColorCube::build_handles(ViewerContext* ctx)
{
    if (ctx->transform_mode() == VIEWER_2D)
        return;

    validate(false);
    build_input_handles(ctx); // inputs are drawn in current world space

    build_knob_handles(ctx);
    add_draw_handle(ctx);
}


void js_ColorCube::draw_handle(ViewerContext* ctx)
{
    Iop * _image = dynamic_cast<Iop*>(Op::input0());
    Iop * _sample = dynamic_cast<Iop*>(Op::input(3));

    if ( _image )
    {
        Hash _testHash = _image->hash();

        if ( _sample )
            _testHash.append( _sample->hash() );

        _testHash.append( m_drawSize );
        _testHash.append( m_centerColor[0] );
        _testHash.append( m_centerColor[1] );
        _testHash.append( m_centerColor[2] );
        _testHash.append( m_doKeyDraw );
        _testHash.append( m_doKey );
        _testHash.append( m_displaySample );
        _testHash.append( m_sampleColor[0] );
        _testHash.append( m_sampleColor[1] );
        _testHash.append( m_sampleColor[2] );

        _testHash.append( m_doDespill );
        _testHash.append( m_doDespill );
        _testHash.append( m_colorBalance );
        _testHash.append( m_colorMult );
        _testHash.append( m_tolMat.a00 );
        _testHash.append( m_tolMat.a01 );
        _testHash.append( m_tolMat.a02 );
        _testHash.append( m_tolMat.a03 );
        _testHash.append( m_tolMat.a10 );
        _testHash.append( m_tolMat.a11 );
        _testHash.append( m_tolMat.a12 );
        _testHash.append( m_tolMat.a13 );
        _testHash.append( m_tolMat.a20 );
        _testHash.append( m_tolMat.a21 );
        _testHash.append( m_tolMat.a22 );
        _testHash.append( m_tolMat.a23 );
        _testHash.append( m_tolMat.a30 );
        _testHash.append( m_tolMat.a31 );
        _testHash.append( m_tolMat.a32 );
        _testHash.append( m_tolMat.a33 );

        _testHash.append( m_softMat.a00 );
        _testHash.append( m_softMat.a01 );
        _testHash.append( m_softMat.a02 );
        _testHash.append( m_softMat.a03 );
        _testHash.append( m_softMat.a10 );
        _testHash.append( m_softMat.a11 );
        _testHash.append( m_softMat.a12 );
        _testHash.append( m_softMat.a13 );
        _testHash.append( m_softMat.a20 );
        _testHash.append( m_softMat.a21 );
        _testHash.append( m_softMat.a22 );
        _testHash.append( m_softMat.a23 );
        _testHash.append( m_softMat.a30 );
        _testHash.append( m_softMat.a31 );
        _testHash.append( m_softMat.a32 );
        _testHash.append( m_softMat.a33 );
        _testHash.append( m_doWarp );

        _testHash.append( m_doKKey );
        _testHash.append( m_kTolMin );
        _testHash.append( m_kTolMax );


        if ( _testHash != m_hash )
        {
            m_hash = _testHash;
            updateColorList();
        }
        
        if ( m_displayListId != -1 )
        {
            glPointSize(m_drawSize);
            glCallList(m_displayListId);
        }
    }
}
void js_ColorCube::knobs(Knob_Callback f)
{
    Tab_knob( f, "3D histogram" );
    Float_knob( f, &m_drawSize, "m_drawSize_knob", "Point size" );
    Newline(f);
    Bool_knob( f, &m_doKeyDraw, "m_doKeyDraw_knob", "Draw Keying result" );
    Newline(f);
    Bool_knob( f, &m_displaySample, "m_displaySample_knob", "Display sample color" );
    Newline(f);
    Color_knob(f, m_sampleColor, DD::Image::IRange(-1, 1), "m_sampleColor_knob", "Sample color");

    Tab_knob( f, "3D keyer" );
    Bool_knob( f, &m_doKey, "m_doKey_knob", "unleash the keyer?" );
    Color_knob(f, m_centerColor, DD::Image::IRange(-1, 1), "m_centerColor_knob", "CenterColor");
    Newline(f);
    /*
    Bool_knob( f, &m_doKKey, "m_doKKey_knob", "Add sample" );
    Float_knob( f, &m_kTolMin, "m_kTolMin_knob", "Sample min" );
    Float_knob( f, &m_kTolMax, "m_kTolMax_knob", "Sample max" );
    */

    Tab_knob( f, "Despill" );
    Bool_knob( f, &m_doDespill, "m_doDespill_knob", "Unleash the despill?" );
    Bool_knob( f, &m_blueScreen, "m_blueScreen_knob", "Bluescreen?" );
    Float_knob( f, &m_colorBalance, "m_colorBalance_knob", "Balance" );
    Float_knob( f, &m_colorMult, "m_colorMult_knob", "Mult" );

    Tab_knob( f, "Color Warper" );
    Bool_knob( f, &m_doWarp, "m_doWarp_knob", "Unleash the warper?" );
    LookupCurves_knob( f, &m_warperCurve, "m_warperCurve_knob", "Warper");
}

static Op* js_ColorCube_constructor(Node *node)
{
    return new js_ColorCube(node);
}

const Iop::Description js_ColorCube::description("js_ColorCube", js_ColorCube_constructor);


